"""
=========================================================
HHL 相位估计 + Pauli Trotter (Suzuki2) 演化 Oracle 模块
=========================================================

【文件用途】
-------------
此模块实现 HHL 算法中所需的 **受控哈密顿演化算子**
    U(τ) = exp(-i * τ * H)
的近似构造方法，其中 H 为可分解为若干 Pauli 张量积的哈密顿量。

特点：
- 支持使用 Pauli 基分解（X、Y、Z、I）对任意 Hermitian 矩阵 H 展开；
- 采用一阶 Trotter 或二阶 Suzuki 分片（Suzuki-Trotter decomposition）近似；
- 使用 cqlib 的原生门实现（如 qc.crz, qc.cx, qc.h, qc.s, qc.sd）；
- 适用于 HHL 算法、量子相位估计（PE）或一般的量子动力学仿真任务。

类型约定：
-----------
EvolutionOracle = Callable[[Circuit, int, Iterable[int], float], None]
表示一个受控时间演化子：
    evolution(qc, ctrl, wrk, tau)
其中：
    - qc: Circuit 对象（cqlib 电路）
    - ctrl: 控制比特（Phase Estimation 计数寄存器位）
    - wrk: 被作用的工作量子比特组（系统寄存器）
    - tau: 演化时间参数 t 或 τ

依赖：
    - circuits.algorithm_circuit.Algorithm_Circuit
    - numpy, math

"""

from __future__ import annotations
from typing import Callable, List, Tuple, Dict, Iterable
import math
import numpy as np

# 导入 cqlib 电路基类
from circuits.algorithm_circuit import Algorithm_Circuit as Circuit


# ============================================================
# 类型定义：EvolutionOracle（受控哈密顿演化算子）
# ============================================================
EvolutionOracle = Callable[[Circuit, int, Iterable[int], float], None]


# ============================================================
# 单比特 Pauli 矩阵表
# ============================================================
PAULI_SINGLE = {
    'I': np.array([[1, 0], [0, 1]], dtype=complex),
    'X': np.array([[0, 1], [1, 0]], dtype=complex),
    'Y': np.array([[0, -1j], [1j, 0]], dtype=complex),
    'Z': np.array([[1, 0], [0, -1]], dtype=complex),
}


# ============================================================
# 工具函数：生成所有 n 比特 Pauli 张量积字符串
# ============================================================
def _all_pauli_strings(n: int) -> List[str]:
    """
    返回长度为 n 的所有 Pauli 字符串组合，例如：
    n=2 → ['II', 'IX', 'IY', 'IZ', 'XI', 'XX', ... 'ZZ']

    规则：
        每个位置 ∈ {'I','X','Y','Z'}，共 4^n 种。
    """
    out = ['I', 'X', 'Y', 'Z']
    for _ in range(n - 1):
        out = [p + q for p in out for q in ['I', 'X', 'Y', 'Z']]
    return out


# ============================================================
# 工具函数：生成多比特 Pauli 矩阵
# ============================================================
def _pauli_matrix(pauli: str) -> np.ndarray:
    """
    构造给定字符串对应的张量积 Pauli 矩阵。
    约定：
        pauli[0] 作用在 wrk[0]，即右侧的 Kronecker 因子（LSB）。
    例如：
        'XZ' → Z ⊗ X （高位左乘）
    """
    M = PAULI_SINGLE[pauli[0]]
    for ch in pauli[1:]:
        M = np.kron(PAULI_SINGLE[ch], M)
    return M


# ============================================================
# 工具函数：把哈密顿矩阵分解为 Pauli 叠加形式
# ============================================================
def decompose_to_paulis_with_identity(H: np.ndarray) -> Tuple[float, int, Dict[str, float]]:
    """
    把任意 Hermitian 矩阵 H 展开为 Pauli 张量积的线性组合：
        H = c_I * I + Σ_p c_p * P_p
    返回：
        c_I   —— 恒等项系数（实数）
        n     —— 量子比特数（log2(dim)）
        coeffs —— 非恒等项的 Pauli 字符串 -> 系数字典
    """
    dim = H.shape[0]
    n = int(round(math.log2(dim)))
    if 2 ** n != dim:
        raise ValueError("H must be 2^n x 2^n.")

    # 恒等项系数（Trace 平均值）
    c_I = float((np.trace(H).real) / (2 ** n))

    coeffs: Dict[str, float] = {}
    for p in _all_pauli_strings(n):
        c = (np.trace(_pauli_matrix(p).conj().T @ H) / (2 ** n)).real
        if p == 'I' * n:
            continue
        if abs(c) > 1e-12:
            coeffs[p] = float(c)

    return c_I, n, coeffs


# ============================================================
# HamiltonianSimulator：通过 Trotter 分片模拟 exp(-i t H)
# ============================================================
class HamiltonianSimulator:
    """
    用 Trotter 或 Suzuki 二阶分片法来近似受控哈密顿演化：
        e^{-i t H} ≈ Π_p e^{-i t c_p P_p}
    并支持外层受控（即 Controlled-U）。

    参数：
        terms      : Dict[str, float]  —— Pauli 字符串及其系数
        method     : {'first', 'suzuki2'}  —— 分片方法
        num_slices : int  —— 分片数量
    """

    def __init__(self, terms: Dict[str, float], method: str = 'suzuki2', num_slices: int = 1):
        if method not in ('first', 'suzuki2'):
            raise ValueError("method must be 'first' or 'suzuki2'")
        self.terms = terms
        self.method = method
        self.num_slices = max(1, int(num_slices))

    # -------------------------------------------
    # (1) 基变换：将 X/Y 基转为 Z 基，便于相位演化
    # -------------------------------------------
    @staticmethod
    def _basis_change(qc: Circuit, qb: int, op: str):
        """
        把单比特 Pauli 运算符 P ∈ {X, Y, Z, I}
        映射到 Z 基上，使后续 e^{-i θ Z} 演化可实现为 RZ 旋转。

        规则：
            I/Z → 无需操作
            X   → H
            Y   → S† + H
        """
        if op in ('I', 'Z'):
            return
        if op == 'X':
            qc.h(qb)
        elif op == 'Y':
            qc.sd(qb)
            qc.h(qb)

    # -------------------------------------------
    # (2) 回滚：恢复原基（Z -> X/Y）
    # -------------------------------------------
    @staticmethod
    def _basis_uncompute(qc: Circuit, qb: int, op: str):
        """
        撤销前面的基变换，使系统恢复原本的算符基。
        规则：
            I/Z → 无需操作
            X   → H
            Y   → H + S
        """
        if op in ('I', 'Z'):
            return
        if op == 'X':
            qc.h(qb)
        elif op == 'Y':
            qc.h(qb)
            qc.s(qb)

    # -------------------------------------------
    # (3) 核心：添加受控的 Pauli 演化 e^{-i θ ⊗ P_i}
    # -------------------------------------------
    @staticmethod
    def _add_ctrl_pauli_evolution(qc: Circuit, ctrl: int, wrk: Iterable[int], pauli: str, theta: float):
        """
        构造受控算子：
            CU = exp(-i * θ * ⊗_i P_i)
        实现方式：
            1. 将所有非 I 的工作位变换到 Z 基；
            2. 用一串 CNOT 把多比特 Z 作用整合到“锚位”量子比特上；
            3. 在“锚位”上施加 CRZ(2θ)（控制比特为 ctrl）；
            4. 撤销 CNOT 链和基变换。

        要求：
            - cqlib 提供 qc.crz(ctrl, tgt, theta)。
            - pauli 是如 'XZI' 形式的字符串。
        """
        ops = list(pauli)
        idx_local = [i for i, op in enumerate(ops) if op != 'I']
        if not idx_local:
            return  # 若该项为纯恒等，则跳过

        wrk_list = list(wrk)

        # 1) 基变换
        for i in idx_local:
            HamiltonianSimulator._basis_change(qc, wrk_list[i], ops[i])

        # 2) 选取最后一个非 I 位作为锚点
        anchor_i = idx_local[-1]
        anchor_q = wrk_list[anchor_i]

        # 3) CNOT 链：把所有非 I 位的 Z 信息聚合到锚点
        for j in idx_local[:-1]:
            qc.cx(wrk_list[j], anchor_q)

        # 4) 在锚点上施加受控相位旋转（控制：ctrl）
        qc.crz(ctrl, anchor_q, 2.0 * theta)

        # 5) 逆向撤销 CNOT 链
        for j in reversed(idx_local[:-1]):
            qc.cx(wrk_list[j], anchor_q)

        # 6) 撤销基变换
        for i in reversed(idx_local):
            HamiltonianSimulator._basis_uncompute(qc, wrk_list[i], ops[i])

    # -------------------------------------------
    # (4) 一阶 Trotter 分片：顺序施加每个 e^{-i c_p P_p dt}
    # -------------------------------------------
    def _slice_first(self, qc: Circuit, ctrl: int, wrk: Iterable[int], dt: float):
        for p, c in self.terms.items():
            self._add_ctrl_pauli_evolution(qc, ctrl, wrk, p, theta=c * dt)

    # -------------------------------------------
    # (5) 二阶 Suzuki 分片：
    #     e^{-i H dt} ≈ Π_p e^{-i c_p P_p dt/2} · Π_p e^{-i c_p P_p dt/2}^†
    # -------------------------------------------
    def _slice_suzuki2(self, qc: Circuit, ctrl: int, wrk: Iterable[int], dt: float):
        half = 0.5 * dt
        items = list(self.terms.items())

        # 正序执行半步
        for p, c in items:
            self._add_ctrl_pauli_evolution(qc, ctrl, wrk, p, theta=c * half)

        # 反序执行半步
        for p, c in reversed(items):
            self._add_ctrl_pauli_evolution(qc, ctrl, wrk, p, theta=c * half)

    # -------------------------------------------
    # (6) 添加完整的受控时间演化 e^{-i τ H}
    # -------------------------------------------
    def add_controlled_time_evolution(self, qc: Circuit, ctrl: int, wrk: Iterable[int], total_time: float):
        """
        将完整的受控 exp(-i τ H) 添加到电路中。

        参数：
            qc         : Circuit
            ctrl       : 控制量子比特
            wrk        : 工作寄存器（作用 H 的对象）
            total_time : 总演化时间 τ

        实现：
            - 把 τ 均分为 num_slices 份；
            - 每份调用 _slice_first 或 _slice_suzuki2；
            - 每份时间步长为 dt = τ / num_slices。
        """
        dt = total_time / self.num_slices
        for _ in range(self.num_slices):
            if self.method == 'first':
                self._slice_first(qc, ctrl, wrk, dt)
            else:
                self._slice_suzuki2(qc, ctrl, wrk, dt)


# ============================================================
# 构造函数：从矩阵 H 生成演化 Oracle
# ============================================================
def build_pauli_evolution_oracle_from_matrix(
    H: np.ndarray,
    method: str = 'suzuki2',
    num_slices: int = 2
) -> EvolutionOracle:
    """
    构造并返回一个 EvolutionOracle，可被 HHL/PE 调用。

    参数：
        H          : Hermitian 矩阵（维度需为 2^n × 2^n）
        method     : 'first' 或 'suzuki2'
        num_slices : 分片数（越大精度越高）

    返回：
        evolution_oracle(qc, ctrl, wrk, tau)
            —— 可直接在电路中调用的函数，实现受控 exp(-i τ H)。

    内部步骤：
    ----------
    1. 用 Pauli 展开分解 H = c_I * I + Σ_p c_p P_p；
    2. 创建 HamiltonianSimulator 实例；
    3. 构造受控演化函数：
           qc.rz(ctrl, -c_I * τ)  # 恒等项相位
           sim.add_controlled_time_evolution(...)

    注：
        - 恒等项 e^{-i c_I τ I} 仅产生整体相位，
          在受控场景中表现为控制位上的 RZ(-c_I * τ)。
    """
    c_I, n, coeffs = decompose_to_paulis_with_identity(H)
    sim = HamiltonianSimulator(coeffs, method=method, num_slices=num_slices)

    def evolution_oracle(qc: Circuit, ctrl: int, wrk: Iterable[int], tau: float):
        if abs(c_I) > 1e-12:
            qc.rz(ctrl, - c_I * tau)
        sim.add_controlled_time_evolution(qc, ctrl, wrk, total_time=tau)

    return evolution_oracle
